#!/usr/bin/env python

from numpy import linspace,zeros
from scipy.linalg import solve
from collocation import lhs, rhs
from bernstein import interpolant
from exact import exact_solution
from matplotlib import pyplot as plt

n = 40  # Bernstein polynomial order

a = 0.  # Solution interval [a,b]
b = 1.

beta = zeros(n+1)  # Bernstein polynomial coefficients - the principal unknown
beta[0] = 0.       # Boundary conditions - direct imposition
beta[n] = 0.

nd = linspace(a,b,n+1) # nodes

f = zeros(n-1)
for i in range(n-1):
	x = nd[i+1]
	f[i] = rhs(x) - beta[0]*lhs(0,n,a,b,x)-beta[n]*lhs(n,n,a,b,x)  # RHS vector entry

K = zeros( (n-1,n-1) )
for i in range(n-1):
	x = nd[i+1]
	for j in range(n-1):
		K[i,j] = lhs(j+1,n,a,b,x)  # LHS matrix entry

beta[1:n] = solve(K,f)  # Solve linear system

u = zeros(n+1)
for i in range(n+1):
	x = nd[i]
	u[i] = interpolant(beta,n,a,b,x)  # Solution - a Bernstein interpolant constructed using beta's
                                      # at whatever points - we use collocation points here.
uex = zeros(n+1)
for i in range(n+1):
	x = nd[i]
	uex[i] = exact_solution(x)

# Print absolute error
print u-uex

plt.plot(nd,u,'r', nd,uex,'o', nd, u-uex, 'g')
plt.legend(('interpolant', 'analytical', 'abs-err'))
plt.show()



